import numpy as np
import open3d as o3d
import random
from time import time

import copy
def readnpy(path, rtype='pcd'):
    points = np.load(path)
    npy = points.astype(np.float32)
    return npy

def npy2pcd(npy, ind=-1):
    colors = [[1.0, 0, 0],
              [0, 1.0, 0],
              [0, 0, 1.0]]
    color = colors[ind] if ind < 3 else [random.random() for _ in range(3)]
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(npy)
    if ind >= 0:
        pcd.paint_uniform_color(color)
    return pcd


def pcd2npy(pcd):
    npy = np.asarray(pcd.points)
    return npy
def crop_pc(downpcd):
    downpcd.estimate_normals(
        search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30))
    aabb = downpcd.get_axis_aligned_bounding_box()
    oriented_bounding_box = o3d.geometry.OrientedBoundingBox.create_from_axis_aligned_bounding_box(aabb)
    point_cloud_crop = downpcd.crop(oriented_bounding_box)
    return point_cloud_crop,aabb

def display_inlier_outlier(cloud, ind):
    inlier_cloud = cloud.select_by_index(ind)
    outlier_cloud = cloud.select_by_index(ind, invert=True)

    print("Showing outliers (red) and inliers (gray): ")
    outlier_cloud.paint_uniform_color([1, 0, 0])
    inlier_cloud.paint_uniform_color([0.8, 0.8, 0.8])
    o3d.visualization.draw_geometries([inlier_cloud, outlier_cloud],
                                      zoom=0.3412,
                                      front=[0.4257, -0.2125, -0.8795],
                                      lookat=[2.6172, 2.0475, 1.532],
                                      up=[-0.0694, -0.9768, 0.2024])
def draw_registration_result(source, target, transformation):
    source_temp = copy.deepcopy(source)
    target_temp = copy.deepcopy(target)
    source_temp.paint_uniform_color([1, 0.706, 0])
    target_temp.paint_uniform_color([0, 0.651, 0.929])
    source_temp.transform(transformation)
    o3d.visualization.draw_geometries([source_temp, target_temp],
                                      zoom=0.4459,
                                      front=[0.9288, -0.2951, -0.2242],
                                      lookat=[1.6784, 2.0612, 1.4451],
                                      up=[-0.3402, -0.9189, -0.1996])
if __name__=='__main__':
    pcd1=readnpy('/home/wsl/gitee/3D/data/dianchi/train_data/92804.npy')
    pcd2=readnpy('/home/wsl/gitee/3D/data/dianchi/train_data/92803.npy')
    # print(pcd2[:, 0].max(), pcd2[:, 0].min(), pcd2[:, 1].max(), pcd2[:, 1].min(), pcd2[:, 2].max(), pcd2[:, 2].min())
    # print(pcd2.shape)
    # median_pcd1=np.median(pcd1[:,1])
    # median_pcd2=np.median(pcd2[:,1])
    # index1=pcd1[:,1]>median_pcd1
    # pcd1=pcd1[index1]
    # index2 = pcd2[:, 1] > median_pcd2
    # pcd2 = pcd2[index2]
    #
    # pcd1=npy2pcd(pcd1,0)
    # pcd2=npy2pcd(pcd2,2)
    # ratio=1
    # pcd1 = pcd1.voxel_down_sample(voxel_size=ratio)
    # pcd2 = pcd2.voxel_down_sample(voxel_size=ratio)
    #
    # print("Apply point-to-plane ICP")
    # threshold = 50
    # trans_init = np.asarray([[1, 0.0, 0.0, 0.0],
    #                          [0.0, 1.0, 0.0, 0.0],
    #                          [0.0, 0.0, 1.0, 0.0],
    #                          [0.0, 0.0, 0.0, 1.0]])
    # pcd1.estimate_normals(
    #     search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30))
    # pcd2.estimate_normals(
    #     search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.1, max_nn=30))
    # reg_p2l = o3d.pipelines.registration.registration_icp(
    #     pcd1, pcd2, threshold, trans_init,
    #     o3d.pipelines.registration.TransformationEstimationPointToPlane(),
    #     o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=200))
    # print(reg_p2l)
    # print("Transformation is:")
    # print(reg_p2l.transformation)
    # draw_registration_result(pcd1,pcd2, reg_p2l.transformation)
    #
    #cl, ind = pcd1.remove_radius_outlier(nb_points=16, radius=0.05)
    #cl, ind = pcd2.remove_radius_outlier(nb_points=16, radius=0.05)
    # cl, ind = pcd1.remove_statistical_outlier(nb_neighbors=10,std_ratio=2.0)
    # cl, ind = pcd2.remove_statistical_outlier(nb_neighbors=10, std_ratio=2.0)
    #display_inlier_outlier(pcd1,ind)
    # pcd1,box1=crop_pc(pcd1)
    # pcd2,box2=crop_pc(pcd2)
    #o3d.visualization.draw_geometries([pcd1,pcd2])